/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Haystack Software Inc. All rights reserved.
 *  Licensed under the PolyForm Strict License 1.0.0. See License.txt in the project root for
 *  license information.
 *--------------------------------------------------------------------------------------------*/

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See code-license.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import { binarySearch, coalesceInPlace, equals } from "vs/base/common/arrays"
import {
  CancellationToken,
  CancellationTokenSource,
} from "vs/base/common/cancellation"
import { onUnexpectedExternalError } from "vs/base/common/errors"
import { Iterable } from "vs/base/common/iterator"
import { LRUCache } from "vs/base/common/map"
import { commonPrefixLength } from "vs/base/common/strings"
import { URI } from "vs/base/common/uri"
import { IPosition, Position } from "vs/editor/common/core/position"
import { IRange, Range } from "vs/editor/common/core/range"
import { ITextModel } from "vs/editor/common/model"
import {
  DocumentSymbol,
  DocumentSymbolProvider,
} from "vs/editor/common/languages"
import { MarkerSeverity } from "vs/platform/markers/common/markers"
import {
  IFeatureDebounceInformation,
  ILanguageFeatureDebounceService,
} from "vs/editor/common/services/languageFeatureDebounce"
import { createDecorator } from "vs/platform/instantiation/common/instantiation"
import {
  InstantiationType,
  registerSingleton,
} from "vs/platform/instantiation/common/extensions"
import { IModelService } from "vs/editor/common/services/model"
import { DisposableStore } from "vs/base/common/lifecycle"
import { LanguageFeatureRegistry } from "vs/editor/common/languageFeatureRegistry"
import { ILanguageFeaturesService } from "vs/editor/common/services/languageFeatures"

export abstract class TreeElement {
  abstract id: string
  abstract children: Map<string, TreeElement>
  abstract parent: TreeElement | undefined

  remove(): void {
    this.parent?.children.delete(this.id)
  }

  static findId(
    candidate: DocumentSymbol | string,
    container: TreeElement
  ): string {
    // complex id-computation which contains the origin/extension,
    // the parent path, and some dedupe logic when names collide
    let candidateId: string
    if (typeof candidate === "string") {
      candidateId = `${container.id}/${candidate}`
    } else {
      candidateId = `${container.id}/${candidate.name}`
      if (container.children.get(candidateId) !== undefined) {
        candidateId = `${container.id}/${candidate.name}_${candidate.range.startLineNumber}_${candidate.range.startColumn}`
      }
    }

    let id = candidateId
    for (let i = 0; container.children.get(id) !== undefined; i++) {
      id = `${candidateId}_${i}`
    }

    return id
  }

  static getElementById(
    id: string,
    element: TreeElement
  ): TreeElement | undefined {
    if (!id) {
      return undefined
    }
    const len = commonPrefixLength(id, element.id)
    if (len === id.length) {
      return element
    }
    if (len < element.id.length) {
      return undefined
    }
    for (const [, child] of element.children) {
      const candidate = TreeElement.getElementById(id, child)
      if (candidate) {
        return candidate
      }
    }
    return undefined
  }

  static size(element: TreeElement): number {
    let res = 1
    for (const [, child] of element.children) {
      res += TreeElement.size(child)
    }
    return res
  }

  static empty(element: TreeElement): boolean {
    return element.children.size === 0
  }
}

export interface IOutlineMarker {
  startLineNumber: number
  startColumn: number
  endLineNumber: number
  endColumn: number
  severity: MarkerSeverity
}

export class OutlineElement extends TreeElement {
  children = new Map<string, OutlineElement>()
  marker: { count: number; topSev: MarkerSeverity } | undefined

  constructor(
    readonly id: string,
    public parent: TreeElement | undefined,
    readonly symbol: DocumentSymbol
  ) {
    super()
  }
}

export class OutlineGroup extends TreeElement {
  children = new Map<string, OutlineElement>()

  constructor(
    readonly id: string,
    public parent: TreeElement | undefined,
    readonly label: string,
    readonly order: number
  ) {
    super()
  }

  getItemEnclosingPosition(position: IPosition): OutlineElement | undefined {
    return position
      ? this._getItemEnclosingPosition(position, this.children)
      : undefined
  }

  private _getItemEnclosingPosition(
    position: IPosition,
    children: Map<string, OutlineElement>
  ): OutlineElement | undefined {
    for (const [, item] of children) {
      if (
        !item.symbol.range ||
        !Range.containsPosition(item.symbol.range, position)
      ) {
        continue
      }
      return this._getItemEnclosingPosition(position, item.children) || item
    }
    return undefined
  }

  updateMarker(marker: IOutlineMarker[]): void {
    for (const [, child] of this.children) {
      this._updateMarker(marker, child)
    }
  }

  private _updateMarker(markers: IOutlineMarker[], item: OutlineElement): void {
    item.marker = undefined

    // find the proper start index to check for item/marker overlap.
    const idx = binarySearch<IRange>(
      markers,
      item.symbol.range,
      Range.compareRangesUsingStarts
    )
    let start: number
    if (idx < 0) {
      start = ~idx
      if (
        start > 0 &&
        Range.areIntersecting(markers[start - 1], item.symbol.range)
      ) {
        start -= 1
      }
    } else {
      start = idx
    }

    const myMarkers: IOutlineMarker[] = []
    let myTopSev: MarkerSeverity | undefined

    for (
      ;
      start < markers.length &&
      Range.areIntersecting(item.symbol.range, markers[start]);
      start++
    ) {
      // remove markers intersecting with this outline element
      // and store them in a 'private' array.
      const marker = markers[start]
      myMarkers.push(marker)
      ;(markers as Array<IOutlineMarker | undefined>)[start] = undefined
      if (!myTopSev || marker.severity > myTopSev) {
        myTopSev = marker.severity
      }
    }

    // Recurse into children and let them match markers that have matched
    // this outline element. This might remove markers from this element and
    // therefore we remember that we have had markers. That allows us to render
    // the dot, saying 'this element has children with markers'
    for (const [, child] of item.children) {
      this._updateMarker(myMarkers, child)
    }

    if (myTopSev) {
      item.marker = {
        count: myMarkers.length,
        topSev: myTopSev,
      }
    }

    coalesceInPlace(markers)
  }
}

export class OutlineModel extends TreeElement {
  static create(
    registry: LanguageFeatureRegistry<DocumentSymbolProvider>,
    textModel: ITextModel,
    token: CancellationToken
  ): Promise<OutlineModel> {
    const cts = new CancellationTokenSource(token)
    const result = new OutlineModel(textModel.uri)
    const provider = registry.ordered(textModel)
    const promises = provider.map((provider, index) => {
      const id = TreeElement.findId(`provider_${index}`, result)
      const group = new OutlineGroup(
        id,
        result,
        provider.displayName ?? "Unknown Outline Provider",
        index
      )

      return Promise.resolve(
        provider.provideDocumentSymbols(textModel, cts.token)
      )
        .then(
          (result) => {
            for (const info of result || []) {
              OutlineModel._makeOutlineElement(info, group)
            }
            return group
          },
          (err) => {
            onUnexpectedExternalError(err)
            return group
          }
        )
        .then((group) => {
          if (!TreeElement.empty(group)) {
            result._groups.set(id, group)
          } else {
            group.remove()
          }
        })
    })

    const listener = registry.onDidChange(() => {
      const newProvider = registry.ordered(textModel)
      if (!equals(newProvider, provider)) {
        cts.cancel()
      }
    })

    return Promise.all(promises)
      .then(() => {
        if (
          cts.token.isCancellationRequested &&
          !token.isCancellationRequested
        ) {
          return OutlineModel.create(registry, textModel, token)
        } else {
          return result._compact()
        }
      })
      .finally(() => {
        cts.dispose()
        listener.dispose()
        cts.dispose()
      })
  }

  private static _makeOutlineElement(
    info: DocumentSymbol,
    container: OutlineGroup | OutlineElement
  ): void {
    const id = TreeElement.findId(info, container)
    const res = new OutlineElement(id, container, info)
    if (info.children) {
      for (const childInfo of info.children) {
        OutlineModel._makeOutlineElement(childInfo, res)
      }
    }
    container.children.set(res.id, res)
  }

  static get(element: TreeElement | undefined): OutlineModel | undefined {
    while (element) {
      if (element instanceof OutlineModel) {
        return element
      }
      element = element.parent
    }
    return undefined
  }

  readonly id = "root"
  readonly parent = undefined

  protected _groups = new Map<string, OutlineGroup>()
  children = new Map<string, OutlineGroup | OutlineElement>()

  protected constructor(readonly uri: URI) {
    super()

    this.id = "root"
    this.parent = undefined
  }

  private _compact(): this {
    let count = 0
    for (const [key, group] of this._groups) {
      if (group.children.size === 0) {
        // empty
        this._groups.delete(key)
      } else {
        count += 1
      }
    }
    if (count !== 1) {
      //
      this.children = this._groups
    } else {
      // adopt all elements of the first group
      const group = Iterable.first(this._groups.values())!
      for (const [, child] of group.children) {
        child.parent = this
        this.children.set(child.id, child)
      }
    }
    return this
  }

  merge(other: OutlineModel): boolean {
    if (this.uri.toString() !== other.uri.toString()) {
      return false
    }
    if (this._groups.size !== other._groups.size) {
      return false
    }
    this._groups = other._groups
    this.children = other.children
    return true
  }

  getItemEnclosingPosition(
    position: IPosition,
    context?: OutlineElement
  ): OutlineElement | undefined {
    let preferredGroup: OutlineGroup | undefined
    if (context) {
      let candidate = context.parent
      while (candidate && !preferredGroup) {
        if (candidate instanceof OutlineGroup) {
          preferredGroup = candidate
        }
        candidate = candidate.parent
      }
    }

    let result: OutlineElement | undefined = undefined
    for (const [, group] of this._groups) {
      result = group.getItemEnclosingPosition(position)
      if (result && (!preferredGroup || preferredGroup === group)) {
        break
      }
    }
    return result
  }

  getItemById(id: string): TreeElement | undefined {
    return TreeElement.getElementById(id, this)
  }

  updateMarker(marker: IOutlineMarker[]): void {
    // sort markers by start range so that we can use
    // outline element starts for quicker look up
    marker.sort(Range.compareRangesUsingStarts)

    for (const [, group] of this._groups) {
      group.updateMarker(marker.slice(0))
    }
  }

  getTopLevelSymbols(): DocumentSymbol[] {
    const roots: DocumentSymbol[] = []
    for (const child of this.children.values()) {
      if (child instanceof OutlineElement) {
        roots.push(child.symbol)
      } else {
        roots.push(
          ...Iterable.map(child.children.values(), (child) => child.symbol)
        )
      }
    }
    return roots.sort((a, b) =>
      Range.compareRangesUsingStarts(a.range, b.range)
    )
  }

  asListOfDocumentSymbols(): DocumentSymbol[] {
    const roots = this.getTopLevelSymbols()
    const bucket: DocumentSymbol[] = []
    OutlineModel._flattenDocumentSymbols(bucket, roots, "")
    return bucket.sort(
      (a, b) =>
        Position.compare(
          Range.getStartPosition(a.range),
          Range.getStartPosition(b.range)
        ) ||
        Position.compare(
          Range.getEndPosition(b.range),
          Range.getEndPosition(a.range)
        )
    )
  }

  private static _flattenDocumentSymbols(
    bucket: DocumentSymbol[],
    entries: DocumentSymbol[],
    overrideContainerLabel: string
  ): void {
    for (const entry of entries) {
      bucket.push({
        kind: entry.kind,
        tags: entry.tags,
        name: entry.name,
        detail: entry.detail,
        containerName: entry.containerName || overrideContainerLabel,
        range: entry.range,
        selectionRange: entry.selectionRange,
        children: undefined, // we flatten it...
      })

      // Recurse over children
      if (entry.children) {
        OutlineModel._flattenDocumentSymbols(bucket, entry.children, entry.name)
      }
    }
  }
}

export const IOutlineModelService = createDecorator<IOutlineModelService>(
  "IOutlineModelService"
)

export interface IOutlineModelService {
  _serviceBrand: undefined
  getOrCreate(
    model: ITextModel,
    token: CancellationToken
  ): Promise<OutlineModel>
  getDebounceValue(textModel: ITextModel): number
}

interface CacheEntry {
  versionId: number
  provider: DocumentSymbolProvider[]

  promiseCnt: number
  source: CancellationTokenSource
  promise: Promise<OutlineModel>
  model: OutlineModel | undefined
}

export class OutlineModelService implements IOutlineModelService {
  declare _serviceBrand: undefined

  private readonly _disposables = new DisposableStore()
  private readonly _debounceInformation: IFeatureDebounceInformation
  private readonly _cache = new LRUCache<string, CacheEntry>(10, 0.7)

  constructor(
    @ILanguageFeaturesService
    private readonly _languageFeaturesService: ILanguageFeaturesService,
    @ILanguageFeatureDebounceService debounces: ILanguageFeatureDebounceService,
    @IModelService modelService: IModelService
  ) {
    this._debounceInformation = debounces.for(
      _languageFeaturesService.documentSymbolProvider,
      "DocumentSymbols",
      { min: 350 }
    )

    // don't cache outline models longer than their text model
    this._disposables.add(
      modelService.onModelRemoved((textModel) => {
        this._cache.delete(textModel.id)
      })
    )
  }

  dispose(): void {
    this._disposables.dispose()
  }

  async getOrCreate(
    textModel: ITextModel,
    token: CancellationToken
  ): Promise<OutlineModel> {
    const registry = this._languageFeaturesService.documentSymbolProvider
    const provider = registry.ordered(textModel)

    let data = this._cache.get(textModel.id)
    if (
      !data ||
      data.versionId !== textModel.getVersionId() ||
      !equals(data.provider, provider)
    ) {
      const source = new CancellationTokenSource()
      data = {
        versionId: textModel.getVersionId(),
        provider,
        promiseCnt: 0,
        source,
        promise: OutlineModel.create(registry, textModel, source.token),
        model: undefined,
      }
      this._cache.set(textModel.id, data)

      const now = Date.now()
      data.promise
        .then((outlineModel) => {
          data!.model = outlineModel
          this._debounceInformation.update(textModel, Date.now() - now)
        })
        .catch((_err) => {
          this._cache.delete(textModel.id)
        })
    }

    if (data.model) {
      // resolved -> return data
      return data.model
    }

    // increase usage counter
    data.promiseCnt += 1

    const listener = token.onCancellationRequested(() => {
      // last -> cancel provider request, remove cached promise
      if (--data.promiseCnt === 0) {
        data.source.cancel()
        this._cache.delete(textModel.id)
      }
    })

    try {
      return await data.promise
    } finally {
      listener.dispose()
    }
  }

  getDebounceValue(textModel: ITextModel): number {
    return this._debounceInformation.get(textModel)
  }
}

registerSingleton(
  IOutlineModelService,
  OutlineModelService,
  InstantiationType.Delayed
)
